package datastructure.book.dp._2_linear_pre_section.single.all;

import org.junit.jupiter.api.Test;

import java.util.*;

public class _18_Solution {

    public int numberOfArithmeticSlices(int[] nums) {
        if (nums == null || nums.length < 3) {
            return 0;
        }
        Arrays.sort(nums);
        HashMap<Integer, List<List<Integer>>> map = new HashMap<>();
        for (int i = 0; i < nums.length; i++) {
            for (int j = i-1; j >= 0; j--) {
                int sub = nums[i] - nums[j];
                List<List<Integer>> lists = map.get(sub);
                if (lists == null) {
                    lists = new ArrayList<>();
                    ArrayList<Integer> subList = new ArrayList<>();
                    subList.add(nums[j]);
                    subList.add(nums[i]);
                    lists.add(subList);
                    map.put(sub, lists);
                } else {
                    boolean find = false;
                    for (List<Integer> list : lists) {
                        int last = list.get(list.size()-1);
                        if (last == nums[j]) {
                            list.add(nums[i]);
                            find = true;
                            break;
                        }
                    }
                    if (!find) {
                        ArrayList<Integer> subList = new ArrayList<>();
                        subList.add(nums[j]);
                        subList.add(nums[i]);
                        lists.add(subList);
                    }
                }
            }
        }
        int cnt = 0;
        for (Map.Entry<Integer, List<List<Integer>>> entry : map.entrySet()) {
            for (List<Integer> integers : entry.getValue()) {
                if (integers.size() < 3) {
                    continue;
                }
                for (int i = 1; i <= integers.size() - 3 + 1; i++) {
                    cnt+=i;
                }
            }
        }
        return cnt;
    }

    public int numberOfArithmeticSlices2(int[] nums){
        if (nums.length < 3) {
            return 0;
        }
        int result = 0;
        int cnt = 0;
        for (int i = 2; i < nums.length; i++) {
            if (nums[i] - nums[i - 1] == nums[i - 1] - nums[i - 2]) {
                cnt++;
                result += cnt;
            } else {
                cnt = 0;
            }
        }
        return result;
    }

    @Test
    public void test(){
        int[] arr = {1,2,3,4,5,6};
//        System.out.println(numberOfArithmeticSlices(arr));
        System.out.println(numberOfArithmeticSlices2(arr));
    }
}
